/*
 * AIEngine a new generation network intrusion detection system.
 *
 * Copyright (C) 2013-2023  Luis Campo Giralte
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301, USA.
 *
 * Written by Luis Campo Giralte <luis.camp0.2009@gmail.com>
 *
 */
#include "BitcoinProtocol.h"
#include <iomanip>

namespace aiengine {

// List of support bitcoin commands
std::unordered_map<std::string, BitcoinCommandType> BitcoinProtocol::commands_ {
	{ "version", 	std::make_tuple(BC_CMD_VERSION,		"version",	0,	20, [](BitcoinInfo &a) {} ) }, // ok
	{ "verack", 	std::make_tuple(BC_CMD_VERACK,		"version ack",	0,	20, [](BitcoinInfo &a) {} ) }, // ok
	{ "addr", 	std::make_tuple(BC_CMD_ADDR,		"network addr",	0,	24, [](BitcoinInfo &a) {} ) }, // ok
	{ "inv", 	std::make_tuple(BC_CMD_INV,		"inv",		0,	20, [](BitcoinInfo &a) {} ) },
	{ "getdata", 	std::make_tuple(BC_CMD_GETDATA,		"getdata",	0,	20, [](BitcoinInfo &a) {} ) },
	{ "notfound", 	std::make_tuple(BC_CMD_NOTFOUND,	"not found",	0,	20, [](BitcoinInfo &a) {} ) },
	{ "getblocks", 	std::make_tuple(BC_CMD_GETBLOCKS,	"get blocks",	0,	24, [](BitcoinInfo &a) {} ) }, // ok
	{ "getheaders", std::make_tuple(BC_CMD_GETHEADERS,	"get headers",	0,	20, [](BitcoinInfo &a) {} ) },
	{ "tx", 	std::make_tuple(BC_CMD_TX,		"transaction",	0,	20, [](BitcoinInfo &a) { a.incTransactions(); } ) },
	{ "block", 	std::make_tuple(BC_CMD_BLOCK,		"block",	0,	24, [](BitcoinInfo &a) { a.incBlocks(); } ) }, // ok
	{ "headers", 	std::make_tuple(BC_CMD_HEADERS,		"headers",	0,	20, [](BitcoinInfo &a) {} ) },
	{ "getaddr", 	std::make_tuple(BC_CMD_GETADDR,		"getaddr",	0,	24, [](BitcoinInfo &a) {} ) }, // ok
	{ "mempool", 	std::make_tuple(BC_CMD_MEMPOOL,		"mempool",	0,	20, [](BitcoinInfo &a) {} ) },
	{ "ping",	std::make_tuple(BC_CMD_PING,		"ping",		0,	20, [](BitcoinInfo &a) {} ) },
	{ "pong",	std::make_tuple(BC_CMD_PONG,		"pong",		0,	20, [](BitcoinInfo &a) {} ) },
	{ "reject",	std::make_tuple(BC_CMD_REJECT,		"reject",	0,	20, [](BitcoinInfo &a) { a.incRejects(); } ) },
	{ "alert",	std::make_tuple(BC_CMD_ALERT,		"alert",	0,	20, [](BitcoinInfo &a) {} ) }
};

bool BitcoinProtocol::check(const Packet &packet) {

	int length = packet.getLength();

	if (length >= header_size) {
		if ((packet.getSourcePort() == 8333)||(packet.getDestinationPort() == 8333)) {
			setHeader(packet.getPayload());
			if (header_->magic == 0xD9B4BEF9) { // Bitcoin magic value 0xf9beb4d9
				++total_valid_packets_;
				return true;
			}
		}
	}
	++total_invalid_packets_;
        return false;
}

void BitcoinProtocol::setDynamicAllocatedMemory(bool value) {

	info_cache_->setDynamicAllocatedMemory(value);
}

bool BitcoinProtocol::isDynamicAllocatedMemory() const {

	return info_cache_->isDynamicAllocatedMemory();
}

uint64_t BitcoinProtocol::getAllocatedMemory() const {

        uint64_t mem = sizeof(BitcoinProtocol);

        mem += info_cache_->getAllocatedMemory();

        return mem;
}

uint64_t BitcoinProtocol::getCurrentUseMemory() const {

	uint64_t mem = sizeof(BitcoinProtocol);

	mem += info_cache_->getCurrentUseMemory();

	return mem;
}

uint64_t BitcoinProtocol::getTotalAllocatedMemory() const {

        return getAllocatedMemory();
}

uint32_t BitcoinProtocol::getTotalCacheMisses() const {

	return info_cache_->getTotalFails();
}

void BitcoinProtocol::releaseCache() {

        if (FlowManagerPtr fm = flow_mng_.lock(); fm) {
                auto ft = fm->getFlowTable();

                std::ostringstream msg;
                msg << "Releasing " << name() << " cache";

                infoMessage(msg.str());

                uint64_t total_bytes_released_by_flows = 0;
                uint32_t release_flows = 0;

                for (auto &flow: ft) {
                        if (SharedPointer<BitcoinInfo> info = flow->getBitcoinInfo(); info) {
                                total_bytes_released_by_flows += sizeof(info);

                                flow->layer7info.reset();
                                ++release_flows;
                                info_cache_->release(info);
                        }
                }

		data_time_ = boost::posix_time::microsec_clock::local_time();

                std::string funit = "Bytes";

                unitConverter(total_bytes_released_by_flows, funit);

                msg.str("");
                msg << "Release " << release_flows << " flows";
                msg << ", flow " << total_bytes_released_by_flows << " " << funit;
                infoMessage(msg.str());
	}
}

void BitcoinProtocol::releaseFlowInfo(Flow *flow) {

	if (auto info = flow->getBitcoinInfo(); info)
		info_cache_->release(info);
}

void BitcoinProtocol::processFlow(Flow *flow) {

	CPUCycle cycles(&total_cpu_cycles_);
	int length = flow->packet->getLength();
	total_bytes_ += length;
	++flow->total_packets_l7;
	++total_packets_;

	current_flow_ = flow;

	SharedPointer<BitcoinInfo> info = flow->getBitcoinInfo();
        if (!info) {
                if (info = info_cache_->acquire(); !info) {
			logFailCache(info_cache_->name(), flow);
			return;
                }
                flow->layer7info = info;
        }

	const uint8_t *payload = flow->packet->getPayload();
	int offset = 0;

	while (offset + header_size < length) {

		setHeader(&payload[offset]);

		// If no magic no packet :)
		if (header_->magic == 0xD9B4BEF9) {
			const char *cmd = reinterpret_cast<const char*>(&header_->command[0]);
			auto it = commands_.find(cmd);
                	if (it != commands_.end()) {
				int32_t *hits = &std::get<2>(it->second);
				short padding = std::get<3>(it->second);
				int32_t payload_len = getPayloadLength();
				auto callback = std::get<4>(it->second);

				callback(*info.get());

				++total_bitcoin_operations_;
				++(*hits);

				offset = (offset + payload_len + padding) ;
			}
		} else {
			break;
		}
        }
}

void BitcoinProtocol::increaseAllocatedMemory(int value) {

        info_cache_->create(value);
}

void BitcoinProtocol::decreaseAllocatedMemory(int value) {

        info_cache_->destroy(value);
}

void BitcoinProtocol::statistics(std::basic_ostream<char> &out, int level, int32_t limit) const {

	std::ios_base::fmtflags f(out.flags());

	showStatisticsHeader(out, level);

	if (level > 3) {
		for (auto &cmd: commands_) {
			const char *label = std::get<1>(cmd.second);
			int32_t hits = std::get<2>(cmd.second);
			out << "\t" << "Total " << label << ":" << std::right << std::setfill(' ') << std::setw(27 - strlen(label)) << hits << std::endl;
		}
	}
	if ((level > 5)and(flow_forwarder_.lock()))
		flow_forwarder_.lock()->statistics(out);
	if (level > 3)
		info_cache_->statistics(out);

	out.flags(f); // Restore out flags
}

void BitcoinProtocol::statistics(Json &out, int level) const {

	showStatisticsHeader(out, level);

	// Variables that depends on the protocol
	if (level > 3) {
		Json j;

		for (auto &cmd: commands_)
			j.emplace(std::get<1>(cmd.second), std::get<2>(cmd.second));

		out["commands"] = j;
	}
}

CounterMap BitcoinProtocol::getCounters() const {
	CounterMap cm;

        cm.addKeyValue("packets", total_packets_);
        cm.addKeyValue("bytes", total_bytes_);

        for (auto &cmd: commands_) {
        	const char *label = std::get<1>(cmd.second);
                int32_t hits = std::get<2>(cmd.second);

		cm.addKeyValue(label, hits);
	}

        return cm;
}

void BitcoinProtocol::resetCounters() {

	reset();

        total_bitcoin_operations_ = 0;
        for (auto &cmd: commands_)
		std::get<2>(cmd.second) = 0;
}

} // namespace aiengine
